"""
限流器 (Rate Limiter)

实现多种限流算法，控制请求频率，防止系统过载。

支持的限流算法：
1. 令牌桶 (Token Bucket) - 平滑突发流量
2. 滑动窗口 (Sliding Window) - 精确限流
3. 固定窗口 (Fixed Window) - 简单高效

作者: RedFire Team
创建日期: 2025-10-06
版本: v1.0
"""

from abc import ABC, abstractmethod
from typing import Optional, Dict, Any
from datetime import datetime, timedelta
from collections import deque
import asyncio
import logging
import time


class RateLimitExceededError(Exception):
    """限流超限异常"""
    
    def __init__(self, message: str, retry_after: Optional[float] = None):
        super().__init__(message)
        self.retry_after = retry_after


class RateLimiter(ABC):
    """
    限流器抽象基类
    
    所有限流器都应该实现此接口。
    """
    
    def __init__(self, name: str, logger: Optional[logging.Logger] = None):
        """
        初始化限流器
        
        Args:
            name: 限流器名称
            logger: 日志记录器
        """
        self.name = name
        self.logger = logger or logging.getLogger(__name__)
        
        # 统计信息
        self.total_requests = 0
        self.allowed_requests = 0
        self.rejected_requests = 0
    
    @abstractmethod
    async def acquire(self, tokens: int = 1) -> bool:
        """
        尝试获取令牌
        
        Args:
            tokens: 需要的令牌数
        
        Returns:
            bool: 是否成功获取
        """
        pass
    
    @abstractmethod
    async def wait_for_token(self, tokens: int = 1, timeout: Optional[float] = None):
        """
        等待获取令牌
        
        Args:
            tokens: 需要的令牌数
            timeout: 超时时间（秒）
        
        Raises:
            RateLimitExceededError: 超时或无法获取令牌
        """
        pass
    
    @abstractmethod
    def get_available_tokens(self) -> int:
        """获取当前可用令牌数"""
        pass
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        rejection_rate = (
            self.rejected_requests / self.total_requests 
            if self.total_requests > 0 else 0.0
        )
        
        return {
            "name": self.name,
            "total_requests": self.total_requests,
            "allowed_requests": self.allowed_requests,
            "rejected_requests": self.rejected_requests,
            "rejection_rate": rejection_rate
        }


class TokenBucketRateLimiter(RateLimiter):
    """
    令牌桶限流器
    
    特点：
    - 平滑处理突发流量
    - 支持预消费（burst）
    - 令牌匀速生成
    
    适用场景：
    - API限流（允许短时突发）
    - 流量整形
    
    使用示例：
    ```python
    # 创建限流器：每秒10个请求，桶容量20
    limiter = TokenBucketRateLimiter(
        name="api_limiter",
        rate=10.0,
        capacity=20
    )
    
    # 尝试获取令牌
    if await limiter.acquire():
        # 执行操作
        result = await some_api_call()
    else:
        # 被限流
        raise RateLimitExceededError("请求过于频繁")
    ```
    """
    
    def __init__(
        self,
        name: str,
        rate: float,
        capacity: Optional[int] = None,
        initial_tokens: Optional[int] = None,
        logger: Optional[logging.Logger] = None
    ):
        """
        初始化令牌桶限流器
        
        Args:
            name: 限流器名称
            rate: 令牌生成速率（个/秒）
            capacity: 桶容量（默认为rate的2倍）
            initial_tokens: 初始令牌数（默认为满桶）
            logger: 日志记录器
        """
        super().__init__(name, logger)
        
        self.rate = rate
        self.capacity = capacity or int(rate * 2)
        self.tokens = initial_tokens if initial_tokens is not None else self.capacity
        self.last_refill = time.time()
        
        # 锁
        self._lock = asyncio.Lock()
        
        self.logger.info(
            f"令牌桶限流器初始化: {name}, "
            f"速率={rate}/s, 容量={self.capacity}"
        )
    
    async def _refill(self):
        """补充令牌"""
        now = time.time()
        elapsed = now - self.last_refill
        
        # 计算应该补充的令牌数
        new_tokens = elapsed * self.rate
        
        if new_tokens > 0:
            self.tokens = min(self.capacity, self.tokens + new_tokens)
            self.last_refill = now
    
    async def acquire(self, tokens: int = 1) -> bool:
        """
        尝试获取令牌
        
        Args:
            tokens: 需要的令牌数
        
        Returns:
            bool: 是否成功获取
        """
        async with self._lock:
            self.total_requests += 1
            await self._refill()
            
            if self.tokens >= tokens:
                self.tokens -= tokens
                self.allowed_requests += 1
                return True
            else:
                self.rejected_requests += 1
                return False
    
    async def wait_for_token(self, tokens: int = 1, timeout: Optional[float] = None):
        """
        等待获取令牌
        
        Args:
            tokens: 需要的令牌数
            timeout: 超时时间（秒）
        
        Raises:
            RateLimitExceededError: 超时或无法获取令牌
        """
        start_time = time.time()
        
        while True:
            if await self.acquire(tokens):
                return
            
            # 检查超时
            if timeout is not None:
                elapsed = time.time() - start_time
                if elapsed >= timeout:
                    raise RateLimitExceededError(
                        f"等待令牌超时: {timeout}s",
                        retry_after=self._calculate_retry_after(tokens)
                    )
            
            # 等待一小段时间
            await asyncio.sleep(0.01)
    
    def _calculate_retry_after(self, tokens: int) -> float:
        """计算需要等待的时间"""
        needed = max(0, tokens - self.tokens)
        return needed / self.rate
    
    def get_available_tokens(self) -> int:
        """获取当前可用令牌数"""
        return int(self.tokens)
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        stats = super().get_stats()
        stats.update({
            "rate": self.rate,
            "capacity": self.capacity,
            "available_tokens": int(self.tokens)
        })
        return stats


class SlidingWindowRateLimiter(RateLimiter):
    """
    滑动窗口限流器
    
    特点：
    - 精确限流，无边界效应
    - 内存占用较高（记录每次请求）
    - 适合低频、精确场景
    
    适用场景：
    - 严格的速率限制
    - 用户级别限流
    
    使用示例：
    ```python
    # 创建限流器：每分钟100个请求
    limiter = SlidingWindowRateLimiter(
        name="user_limiter",
        max_requests=100,
        window_seconds=60
    )
    
    # 尝试获取
    if await limiter.acquire():
        # 执行操作
        pass
    ```
    """
    
    def __init__(
        self,
        name: str,
        max_requests: int,
        window_seconds: float,
        logger: Optional[logging.Logger] = None
    ):
        """
        初始化滑动窗口限流器
        
        Args:
            name: 限流器名称
            max_requests: 窗口内最大请求数
            window_seconds: 窗口大小（秒）
            logger: 日志记录器
        """
        super().__init__(name, logger)
        
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        
        # 请求时间戳队列
        self.request_times: deque = deque()
        
        # 锁
        self._lock = asyncio.Lock()
        
        self.logger.info(
            f"滑动窗口限流器初始化: {name}, "
            f"限制={max_requests}/{window_seconds}s"
        )
    
    async def _cleanup_old_requests(self):
        """清理过期的请求记录"""
        now = time.time()
        cutoff = now - self.window_seconds
        
        while self.request_times and self.request_times[0] < cutoff:
            self.request_times.popleft()
    
    async def acquire(self, tokens: int = 1) -> bool:
        """
        尝试获取令牌
        
        Args:
            tokens: 需要的令牌数（这里表示请求数）
        
        Returns:
            bool: 是否成功获取
        """
        async with self._lock:
            self.total_requests += 1
            await self._cleanup_old_requests()
            
            # 检查是否超限
            if len(self.request_times) + tokens <= self.max_requests:
                # 记录请求时间
                now = time.time()
                for _ in range(tokens):
                    self.request_times.append(now)
                
                self.allowed_requests += 1
                return True
            else:
                self.rejected_requests += 1
                return False
    
    async def wait_for_token(self, tokens: int = 1, timeout: Optional[float] = None):
        """
        等待获取令牌
        
        Args:
            tokens: 需要的令牌数
            timeout: 超时时间（秒）
        
        Raises:
            RateLimitExceededError: 超时或无法获取令牌
        """
        start_time = time.time()
        
        while True:
            if await self.acquire(tokens):
                return
            
            # 检查超时
            if timeout is not None:
                elapsed = time.time() - start_time
                if elapsed >= timeout:
                    raise RateLimitExceededError(
                        f"等待令牌超时: {timeout}s",
                        retry_after=self._calculate_retry_after()
                    )
            
            # 等待
            await asyncio.sleep(0.1)
    
    def _calculate_retry_after(self) -> float:
        """计算需要等待的时间"""
        if not self.request_times:
            return 0.0
        
        oldest = self.request_times[0]
        now = time.time()
        elapsed = now - oldest
        
        return max(0.0, self.window_seconds - elapsed)
    
    def get_available_tokens(self) -> int:
        """获取当前可用令牌数"""
        return max(0, self.max_requests - len(self.request_times))
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        stats = super().get_stats()
        stats.update({
            "max_requests": self.max_requests,
            "window_seconds": self.window_seconds,
            "current_requests": len(self.request_times),
            "available_slots": self.get_available_tokens()
        })
        return stats


class FixedWindowRateLimiter(RateLimiter):
    """
    固定窗口限流器
    
    特点：
    - 实现简单，性能高
    - 存在边界效应
    - 内存占用低
    
    适用场景：
    - 粗粒度限流
    - 高性能要求
    
    注意：存在边界问题，窗口交界处可能达到2倍限制
    """
    
    def __init__(
        self,
        name: str,
        max_requests: int,
        window_seconds: float,
        logger: Optional[logging.Logger] = None
    ):
        """
        初始化固定窗口限流器
        
        Args:
            name: 限流器名称
            max_requests: 窗口内最大请求数
            window_seconds: 窗口大小（秒）
            logger: 日志记录器
        """
        super().__init__(name, logger)
        
        self.max_requests = max_requests
        self.window_seconds = window_seconds
        
        # 当前窗口
        self.window_start = time.time()
        self.current_requests = 0
        
        # 锁
        self._lock = asyncio.Lock()
        
        self.logger.info(
            f"固定窗口限流器初始化: {name}, "
            f"限制={max_requests}/{window_seconds}s"
        )
    
    async def _check_window(self):
        """检查并重置窗口"""
        now = time.time()
        if now - self.window_start >= self.window_seconds:
            self.window_start = now
            self.current_requests = 0
    
    async def acquire(self, tokens: int = 1) -> bool:
        """
        尝试获取令牌
        
        Args:
            tokens: 需要的令牌数
        
        Returns:
            bool: 是否成功获取
        """
        async with self._lock:
            self.total_requests += 1
            await self._check_window()
            
            if self.current_requests + tokens <= self.max_requests:
                self.current_requests += tokens
                self.allowed_requests += 1
                return True
            else:
                self.rejected_requests += 1
                return False
    
    async def wait_for_token(self, tokens: int = 1, timeout: Optional[float] = None):
        """
        等待获取令牌
        
        Args:
            tokens: 需要的令牌数
            timeout: 超时时间（秒）
        
        Raises:
            RateLimitExceededError: 超时或无法获取令牌
        """
        start_time = time.time()
        
        while True:
            if await self.acquire(tokens):
                return
            
            # 检查超时
            if timeout is not None:
                elapsed = time.time() - start_time
                if elapsed >= timeout:
                    raise RateLimitExceededError(
                        f"等待令牌超时: {timeout}s",
                        retry_after=self._calculate_retry_after()
                    )
            
            # 等待到下一个窗口
            await asyncio.sleep(0.1)
    
    def _calculate_retry_after(self) -> float:
        """计算需要等待的时间"""
        now = time.time()
        elapsed = now - self.window_start
        return max(0.0, self.window_seconds - elapsed)
    
    def get_available_tokens(self) -> int:
        """获取当前可用令牌数"""
        return max(0, self.max_requests - self.current_requests)
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        stats = super().get_stats()
        stats.update({
            "max_requests": self.max_requests,
            "window_seconds": self.window_seconds,
            "current_requests": self.current_requests,
            "available_slots": self.get_available_tokens(),
            "window_start": datetime.fromtimestamp(self.window_start).isoformat()
        })
        return stats

